import json
import numpy as np
import open3d as o3d
from PIL import Image
from scipy.signal import correlate2d
from scipy.ndimage import shift
from utils.misc.post_proc import np_coor2xy, np_coorx2u, np_coory2v
from utils.eval_general import layout_2_depth

def export_3d_model(img_path, layout_path, out_path=None, ignore_floor=False, ignore_ceiling=False, ignore_wall=False):
    """
    从给定的全景图像和房间布局文件生成并导出3D模型

    参数:
    - img_path: 全景图像路径
    - layout_path: 布局文件路径（支持txt或json格式）
    - out_path: 输出的3D模型文件路径 (.ply 格式)
    - ignore_floor: 是否忽略地板渲染
    - ignore_ceiling: 是否忽略天花板渲染
    - ignore_wall: 是否忽略墙面渲染
    """
    # 读取全景图像
    equirect_texture = np.array(Image.open(img_path))
    H, W = equirect_texture.shape[:2]

    # 读取布局信息
    if layout_path.endswith('json'):
        with open(layout_path) as f:
            inferenced_result = json.load(f)
        cor_id = np.array(inferenced_result['uv'], np.float32)
        cor_id[:, 0] *= W
        cor_id[:, 1] *= H
    else:
        cor_id = np.loadtxt(layout_path).astype(np.float32)

    # 将角点转换为深度图
    depth, floor_mask, ceil_mask, wall_mask = layout_2_depth(cor_id, H, W, return_mask=True)
    coorx, coory = np.meshgrid(np.arange(W), np.arange(H))
    us = np_coorx2u(coorx, W)
    vs = np_coory2v(coory, H)
    zs = depth * np.sin(vs)
    cs = depth * np.cos(vs)
    xs = cs * np.sin(us)
    ys = -cs * np.cos(us)

    # 合并遮罩
    mask = np.ones_like(floor_mask)
    if ignore_floor:
        mask &= ~floor_mask
    if ignore_ceiling:
        mask &= ~ceil_mask
    if ignore_wall:
        mask &= ~wall_mask

    # 准备点云和面片
    xyzrgb = np.concatenate([xs[..., None], ys[..., None], zs[..., None], equirect_texture], -1)
    xyzrgb = np.concatenate([xyzrgb, xyzrgb[:, [0]]], 1)
    mask = np.concatenate([mask, mask[:, [0]]], 1)

    lo_tri_template = np.array([[0, 0, 0], [0, 1, 0], [0, 1, 1]])
    up_tri_template = np.array([[0, 0, 0], [0, 1, 1], [0, 0, 1]])
    ma_tri_template = np.array([[0, 0, 0], [0, 1, 1], [0, 1, 0]])

    lo_mask = (correlate2d(mask, lo_tri_template, mode='same') == 3)
    up_mask = (correlate2d(mask, up_tri_template, mode='same') == 3)
    ma_mask = (correlate2d(mask, ma_tri_template, mode='same') == 3) & (~lo_mask) & (~up_mask)
    ref_mask = (
        lo_mask | (correlate2d(lo_mask, np.flip(lo_tri_template, (0, 1)), mode='same') > 0) |
        up_mask | (correlate2d(up_mask, np.flip(up_tri_template, (0, 1)), mode='same') > 0) |
        ma_mask | (correlate2d(ma_mask, np.flip(ma_tri_template, (0, 1)), mode='same') > 0)
    )
    points = xyzrgb[ref_mask]

    ref_id = np.full(ref_mask.shape, -1, np.int32)
    ref_id[ref_mask] = np.arange(ref_mask.sum())

    faces_lo_tri = np.stack([
        ref_id[lo_mask],
        ref_id[shift(lo_mask, [1, 0], cval=False, order=0)],
        ref_id[shift(lo_mask, [1, 1], cval=False, order=0)],
    ], 1)
    faces_up_tri = np.stack([
        ref_id[up_mask],
        ref_id[shift(up_mask, [1, 1], cval=False, order=0)],
        ref_id[shift(up_mask, [0, 1], cval=False, order=0)],
    ], 1)
    faces_ma_tri = np.stack([
        ref_id[ma_mask],
        ref_id[shift(ma_mask, [1, 0], cval=False, order=0)],
        ref_id[shift(ma_mask, [0, 1], cval=False, order=0)],
    ], 1)

    faces = np.concatenate([faces_lo_tri, faces_up_tri, faces_ma_tri])

    # 导出为PLY文件
    if out_path:
        ply_header = '\n'.join([
            'ply',
            'format ascii 1.0',
            f'element vertex {len(points):d}',
            'property float x',
            'property float y',
            'property float z',
            'property uchar red',
            'property uchar green',
            'property uchar blue',
            f'element face {len(faces):d}',
            'property list uchar int vertex_indices',
            'end_header',
        ])
        with open(out_path, 'w') as f:
            f.write(ply_header)
            f.write('\n')
            for x, y, z, r, g, b in points:
                f.write(f'{x:.2f} {y:.2f} {z:.2f} {r:.0f} {g:.0f} {b:.0f}\n')
            for i, j, k in faces:
                f.write(f'3 {i:d} {j:d} {k:d}\n')

    return out_path


